from sklearn.metrics import roc_curve, roc_auc_score
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import scipy
import shelve


class img_tool():
    def __init__(self) -> None:
        pass

    #   这里用来绘制roc曲线图
    #   感觉是需要保存预测值和实际值即可
    # ------------------------------ auc ---------------------------------------------------------
    def draw_roc(self):
        modelName = ["tlsh", "functionSim", "img", "MGMN", "siamese_graphsage"]
        plt.figure(figsize=(7, 7))
        for i in range(len(modelName)):
            print("生成{}模型的roc曲线".format(modelName[i]))
            with shelve.open("/home/chenyongwei/function_sim_project/all_data/indicators/auc/{}_auc".format(modelName[i])) as file:
                pred = file["pred"]
                tar = file["true"]
            fpr, tpr, thresholds = roc_curve(tar, pred)
            auc = roc_auc_score(tar, pred)
            plt.plot(fpr, tpr, label='{} auc: {:.2f}'.format(
                modelName[i], auc))
        plt.plot([0, 1], [0, 1], 'k--')  # 绘制对角线
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('ROC')
        plt.legend(loc='lower right')
        plt.savefig(
            '/home/chenyongwei/function_sim_project/all_data/img/roc_curve.png', dpi=1000)
        plt.show()

    def get_auc_data(self):
        modelName = ["tlsh", "img", "functionSim",
                     "mgmn", "mgmn_without_matching"]
        res = []
        tar = []
        for i in modelName:
            file = shelve.open(
                "/home/chenyongwei/project/malSimModels/dataInf/auc/{}_auc".format(i))
            res.append(np.array(file["pred"]))
            tar.append(np.array(file["actual"]))
            file.close()
        return res, tar

    #   这个先不管吧，感觉没什么用
    #   绘制模型的下降曲线,
    # ------------------------------ loss curve ---------------------------------------------------------
    def get_loss_data(self, a):
        res = []
        with open(a) as file:
            data = file.readlines()
            for i in range(len(data)):
                end = data[i].index(",")
                res.append(float(data[i][7:end]))
        return res

    def draw_loss(self):
        paths = ["/home/chenyongwei/project/malSimModels/dataInf/loss/functionSim_loss.txt",
                 "/home/chenyongwei/project/malSimModels/dataInf/loss/img_loss.txt",
                 "/home/chenyongwei/project/malSimModels/dataInf/loss/mgmn_loss.txt",
                 "/home/chenyongwei/project/malSimModels/dataInf/loss/mgmn_without_matching_loss.txt"]
        names = ["functionSim", "img", "MGMN", "MGMG_without_match"]

        plt.figure(figsize=(20, 4))
        for i in range(len(paths)):
            data = self.get_loss_data(paths[i])
            x = np.array(list(range(len(data))))
            y = np.array(data)
            plt.subplot(1, 4, i+1)
            y_smooth = scipy.signal.savgol_filter(y, 30, 5)
            plt.plot(x, y, alpha=0.4, linewidth=2.5,
                     label="{} Original curve".format(names[i]))
            plt.plot(x, y_smooth, alpha=1,
                     label="{} fit curve".format(names[i]))
            plt.legend(loc='best')
            plt.title("{} loss".format(
                names[i]), fontsize=10, color='blue', fontweight="heavy", loc='left')
        plt.savefig("/home/chenyongwei/project/malSimModels/modelImg/all_loss")

    def draw_confusion_matrix(self, name, matrix, lables):
        # 创建一个 Matplotlib 图表
        plt.figure(figsize=(20, 10), dpi=600)
        # 创建表格并添加数据
        table = plt.table(cellText=matrix, loc='center')
        # 设置表格的样式
        table.auto_set_font_size(False)
        table.set_fontsize(6)
        plt.tight_layout()
        plt.axis('off')

        table_props = table.properties()['celld']

        # # 设置第一行的字体样式
        # for cell in table_props.values():
        #     cell.set_fontsize(2)
        #     cell.get_text().set_weight('bold')
        #     cell.set_facecolor('#C3E6E3')

        # 设置第一列的字体样式
        for i in range(1, len(matrix)):
            cell = table_props[(i, 0)]
            cell.set_fontsize(2)
            cell = table_props[(0, i)]
            cell.set_fontsize(2)

        # 自定义颜色映射
        cmap = mcolors.LinearSegmentedColormap.from_list(
            "", ["white", "darkblue"])
        # 设置单元格颜色
        newMatrix = [row[1:] for row in matrix[1:]]
        maxvalue = np.max(newMatrix)
        for i in range(1, len(matrix)):
            for j in range(1, len(matrix[0])):
                value = matrix[i][j] + 150 if matrix[i][j] != 0 else 0  # 数据值
                value = min(value, maxvalue)
                cell = table[(i, j)]  # 获取单元格
                cell.set_facecolor(cmap(value / maxvalue))  # 根据数据值设置颜色
                if i == j:
                    cell.set_facecolor('#C3E6C3')

        plt.savefig(
            "/home/chenyongwei/function_sim_project/all_data/img/{}_cinfusion_matrix.png".format(name))


if __name__ == "__main__":
    gene_img = img_tool()
    # gene_img.draw_roc()
    gene_img.draw_confusion_matrix(0, 0, 0)
